import sys 
sys.path.append('code')
from feature_eng.feature_fundmantal.alpha_value import Alpha_other, Alphas_value
import os  
from joblib import Parallel, delayed
import itertools
import pandas as pd
import time
import numpy as np

class Value_main():
    
    def __init__(self, columns_root=None, save_root=None, price_root=None, value_root=None) -> None:
        self.price_root = r'data/stock_data/daily/price' if price_root is None else  price_root
        self.value_root = r'data/stock_data/daily/market_value' if value_root is None else  value_root
        self.columns_root = r'data/stock_data/daily/consentrate_price/close.pkl.gzip' if columns_root is None else columns_root        
        self.save_root = r'data/cmodty/my_feature/feature_data_fd' if save_root is None else save_root     
        if not os.path.exists(self.save_root):
            os.makedirs(self.save_root)
        self.column = pd.read_pickle(self.columns_root).columns
            
    def calc_alpha(self, func, name='alpha_v0', *params):
        t1 = time.time()
        lst = []
        for code in self.column:
            price_file = os.path.join(self.price_root,f'{code}.pkl.gzip')
            value_file = os.path.join(self.value_root,f'{code}.pkl.gzip')
            price = pd.read_pickle(price_file).astype(np.float32)
            value = pd.read_pickle(value_file).astype(np.float32)
            df_i = pd.concat([price,value], axis=1)
            df_i = df_i
            alpha_i = pd.DataFrame(index = df_i.index) 
            ary = func(name, df_i, params)
            alpha_i[code] = ary
            lst.append(alpha_i)
        df = pd.concat(lst, axis=1)
        df.columns = self.column
        df = df[df.index>'2005-01-01']
        save_file = os.path.join(self.save_root, f'{name}_{params}.pkl.gzip')
        df.to_pickle(save_file)
        print(f"{name}_{params} calculating cost time:{(time.time()-t1):.2f}s")

    def eval_alpha_v0(self, name, df_i, params):
        alp1 = Alphas_value(df_i)
        ary = eval(f'alp1.{name}')()
        return ary

    def eval_alpha_v1(self, name, df_i, params):
        alp1 = Alphas_value(df_i)
        ary = eval(f'alp1.{name}')(params[0])
        return ary

    def eval_alpha_v2(self, name, df_i, params):
        alp1 = Alphas_value(df_i)
        ary = eval(f'alp1.{name}')(params[0], params[1])
        return ary

    def eval_alpha_v3(self, name, df_i, params):
        alp1 = Alphas_value(df_i)
        ary = eval(f'alp1.{name}')(params[0], params[1], params[2])
        return ary

    def eval_alpha_o0(self, name, df_i, params):
        alp1 = Alpha_other(df_i)
        ary = eval(f'alp1.{name}')()
        return ary

    def eval_alpha_o1(self, name, df_i, params):
        alp1 = Alpha_other(df_i)
        ary = eval(f'alp1.{name}')(params[0])
        return ary

    def eval_alpha_o2(self, name, df_i, params):
        alp1 = Alpha_other(df_i)
        ary = eval(f'alp1.{name}')(params[0], params[1])
        return ary

    def eval_alpha_o3(self, name, df_i, params):
        alp1 = Alpha_other(df_i)
        ary = eval(f'alp1.{name}')(params[0], params[1], params[2])
        return ary
    


def calc_main_fundmental():
    #v0: alpha_v1 ~ v4
    #v1: alpha_v5 ~ v8
    #v2: alpha_v9 ~ v14
    #v3: alpha_v15 ~ v23

    #o0: alpha_o1 ~ o5
    #o1: alpha_o6 ~ o42
    #o2: alpha_o43 ~ o50
    #o3: alpha_o51 ~ o63
    lst_v1 = list(range(5,13))
    lst_v2 = list(range(13,19))
    lst_v3 = list(range(19,28))

    lst_o1 = list(range(6,43))
    lst_o2 = list(range(43,51))
    lst_o3 = list(range(51,64))
    param_lst = [5,10,30, 60, 120]
    ratio_lst = [2,3,4]
    C = Value_main()
    Parallel(n_jobs=16)(delayed(C.calc_alpha)(C.eval_alpha_v0, f'alpha_v{i}', 0) for i in range(1,5))
    Parallel(n_jobs=16)(delayed(C.calc_alpha)(C.eval_alpha_v1, f'alpha_v{i}', j) for i,j in itertools.product(lst_v1, param_lst))
    Parallel(n_jobs=16)(delayed(C.calc_alpha)(C.eval_alpha_v2, f'alpha_v{i}', j, k) for i,j, k in itertools.product(lst_v2, param_lst[:4], ratio_lst))
    Parallel(n_jobs=16)(delayed(C.calc_alpha)(C.eval_alpha_v3, f'alpha_v{i}', j, k, w) for i,j,k,w in itertools.product(lst_v3, param_lst[:4], param_lst[:-1], param_lst[:-1]))

    Parallel(n_jobs=16)(delayed(C.calc_alpha)(C.eval_alpha_o0, f'alpha_o{i}', 0) for i in range(1,6))
    Parallel(n_jobs=16)(delayed(C.calc_alpha)(C.eval_alpha_o1, f'alpha_o{i}', j) for i,j in itertools.product(lst_o1, param_lst))
    Parallel(n_jobs=16)(delayed(C.calc_alpha)(C.eval_alpha_o2, f'alpha_o{i}', j, k) for i,j, k in itertools.product(lst_o2, param_lst[:4], ratio_lst))
    Parallel(n_jobs=16)(delayed(C.calc_alpha)(C.eval_alpha_o3, f'alpha_o{i}', j, k, w) for i,j,k,w in itertools.product(lst_o3, param_lst[:4], param_lst[:-1], param_lst[:-1]))

# calc_main_fundmental()